import pandas as pd
import os
import redis

LOCAL_RAW_DIR = os.path.join(os.path.dirname(__file__), 'raw')
REDIS_HOST = 'localhost'
REDIS_PORT = 6379

def merge_and_save_redis():
    dfs = []
    for file in os.listdir(LOCAL_RAW_DIR):
        if file.endswith('.csv'):
            file_path = os.path.join(LOCAL_RAW_DIR, file)
            # 修改读取方式，原始数据格式为：日期,国家代码,出口额,进口额
            df = pd.read_csv(file_path, header=None, names=['日期', '国家标签', '出口额', '进口额'])
            dfs.append(df)
    
    if not dfs:
        raise ValueError("未找到任何数据文件，请检查data/raw目录")
    
    table1 = pd.concat(dfs, ignore_index=True)
    # 确保日期格式正确
    table1['日期'] = table1['日期'].astype(str)
    table1 = table1.sort_values(['日期', '国家标签']).reset_index(drop=True)
    
    table2 = table1.groupby('日期').agg({
        '出口额': 'sum',
        '进口额': 'sum'
    }).reset_index().sort_values('日期')
    
    r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True)
    r.set('japan_trade_table1', table1.to_json(orient='records'))
    r.set('japan_trade_table2', table2.to_json(orient='records'))
    print("数据合并完成，已存储到Redis")

if __name__ == '__main__':
    merge_and_save_redis()